# -*- coding: utf-8 -*-
"""
@file: baiduOCR
@author: Memory
@date: 2020/10/11
@description: 百度OCR的接口
"""
import re
import time
from PIL import Image
from aip import AipOcr
from taobao_check.tools import pil2bytes, merge_image


class BaiduOCR(object):
    """
    百度OCR
    """

    def __init__(self, app_id: str, app_key: str, app_secret: str):
        """
        以下三个参数需要从百度AI开放平台上获取：https://ai.baidu.com/
        Args:
            app_id:
            app_key:
            app_secret:
        """
        self.interval = 1  # 用于控制两次调用之间的时间间隔
        self.last_call = time.time()  # 上一次的调用时间
        self.aip = AipOcr(app_id, app_key, app_secret)  # 百度官方提供的OCR接口

    def ocr(self, image: Image, verification_code=False) -> list:
        """
        识别一张图像
        Args:
            image: 可以时 PIL.image 对象，也可以是
        Returns:
            如果识别成功，则返回识别的文本列表，识别识别则返回空列表
        """
        if not image:
            return []

        if verification_code:
            image = filter_bg(image).convert("1")
            image = depoint(depoint(image))
            clear_noise(image, 120, 3, 2)

        image = pil2bytes(image)

        if time.time() - self.last_call < self.interval:
            time.sleep(self.last_call + self.interval - self.last_call + 0.1)

        # if verification_code:
        #     dic_result = self.aip.handwriting(image)
        # else:
        dic_result = self.aip.basicAccurate(image)
        self.last_call = time.time()

        if 'error_code' in dic_result:
            print(dic_result)
            return []

        if 'words_result' not in dic_result:
            print(dic_result)
            return []

        res = dic_result['words_result']
        if verification_code:
            res = "".join([words['words'] for words in res if words['words'].strip()])
            print("识别结果:", res)
            res = res.replace(" ", "")
            res = re.sub("[^A-Za-z0-9]+", "", res)
            res = res.lower()
        else:
            res = [words['words'] for words in res]

        return res


def filter_bg(img):
    img = img.convert("RGB")
    w, h = img.size
    pixdata = img.load()
    for x in range(w):
        for y in range(h):
            r, g, b = pixdata[x, y]
            if max(r, g, b) - min(r, g, b) < 18:
                pixdata[x, y] = (255, 255, 255)
            else:
                pixdata[x, y] = (0, 0, 0)
    return img


def depoint(img):  # input: gray image
    pixdata = img.load()
    w, h = img.size
    for y in range(1, h - 1):
        for x in range(1, w - 1):
            count = 0
            if pixdata[x, y - 1] > 245:
                count = count + 1
            if pixdata[x, y + 1] > 245:
                count = count + 1
            if pixdata[x - 1, y] > 245:
                count = count + 1
            if pixdata[x + 1, y] > 245:
                count = count + 1
            if count > 2:
                pixdata[x, y] = 255
    return img


def get_pixel(image, x, y, G, N):
    # 获取像素值
    L = image.getpixel((x, y))

    # 与阈值比较
    if L > G:
        L = True
    else:
        L = False

    nearDots = 0

    if L == (image.getpixel((x - 1, y - 1)) > G):
        nearDots += 1
    if L == (image.getpixel((x - 1, y)) > G):
        nearDots += 1
    if L == (image.getpixel((x - 1, y + 1)) > G):
        nearDots += 1
    if L == (image.getpixel((x, y - 1)) > G):
        nearDots += 1
    if L == (image.getpixel((x, y + 1)) > G):
        nearDots += 1
    if L == (image.getpixel((x + 1, y - 1)) > G):
        nearDots += 1
    if L == (image.getpixel((x + 1, y)) > G):
        nearDots += 1
    if L == (image.getpixel((x + 1, y + 1)) > G):
        nearDots += 1

    if nearDots < N:
        return image.getpixel((x, y - 1))
    else:
        return None


def clear_noise(image, G, N, Z):
    w, h = image.size

    for i in range(0, Z):
        for x in range(1, w - 1):
            for y in range(1, h - 1):
                color = get_pixel(image, x, y, G, N)
                if color is not None:
                    image.putpixel((x, y), color)


app_id = "24476366"
key = "TLWm2bkGv2C32uU7CbRDEeFI"
secret = "K41AGQgeT42erkWIQXGXBoAjoEDQPewA"

OCR = BaiduOCR(app_id, key, secret)

if __name__ == "__main__":
    from pathlib import Path
    from random import choice

    imgs = list(Path("veri_code").iterdir())
    ocr = BaiduOCR(app_id, key, secret)
    image = Image.open(choice(imgs))
    # image = filter_bg(img).convert("1")
    # image = depoint(depoint(image))
    # clear_noise(image, 120, 3, 2)

    print("识别结果:", ocr.ocr(image, True))
    image.show()
    # for idx, im in enumerate(imgs):
    #     print(idx)
    #     img = Image.open(im)
    #     image = filter_bg(img).convert("1")
    #     image = depoint(depoint(image))
    #     # clear_noise(image, 120, 3, 2)
    #     merge_image(img, image).save(f"merge\\{idx:0>4}.png")

